import pybullet #这里导入pybullet库
import pybullet_data
import time
import numpy as np
# pybullet.connect(pybullet.DIRECT)#建立连接，这一步是pybullet所必须的
rendering = True
if not rendering:
    pybullet.connect(pybullet.DIRECT)
else:
    pybullet.connect(pybullet.GUI_SERVER)
pybullet.setAdditionalSearchPath(pybullet_data.getDataPath())#添加pybullet内部的模型路径
debug=False
if not debug:
    pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_GUI, 0)
pybullet.resetDebugVisualizerCamera(cameraDistance=0.3, cameraYaw=60, cameraPitch=-50,
                                    cameraTargetPosition=[0.5, -0.05, 0.1])
pybullet.resetDebugVisualizerCamera(cameraDistance=0.3, cameraYaw=60, cameraPitch=-50,
                                    cameraTargetPosition=[0.5, -0.05, 0.1])
pybullet.setPhysicsEngineParameter(enableFileCaching=0)
pybullet.configureDebugVisualizer(pybullet.COV_ENABLE_SHADOWS, True)
pybullet.setGravity(0, 0, -9.80)#设置重力
pybullet.resetSimulation()
pybullet.getCameraImage(320,200, renderer=pybullet.ER_BULLET_HARDWARE_OPENGL )#得到图像
# pybullet.setRealTimeSimulation(1)

timeStep=1./120.#240.
pybullet.setTimeStep(timeStep)
urdf_paths = 'rm_65_description/urdf/rm_65_rev.urdf'
# pybullet_robot = pybullet.loadURDF(urdf_paths, useFixedBase=False, flags=pybullet.URDF_USE_SELF_COLLISION)
pybullet_robot = pybullet.loadURDF(urdf_paths, basePosition=[0, 0, 0], baseOrientation=pybullet.getQuaternionFromEuler([0., 0., 0.]), flags=pybullet.URDF_USE_SELF_COLLISION)

# robot_position = [0,0,0]
# robot_orn = [0, 0, -1, 1]
# pybullet.resetBasePositionAnd
# pybullet.loadURDF("plane.urdf") 
num_joints = pybullet.getNumJoints(pybullet_robot) # 
# 下面循环主要用来获取机械臂转动关节数量
jointIds_c = []
for j in range (num_joints):
     pybullet.changeDynamics(pybullet_robot, j, linearDamping=0, angularDamping=0)
     info_c = pybullet.getJointInfo(pybullet_robot, j)
     # print(info)
     jointName_c = info_c[1]
     jointType_c = info_c[2]
     if (jointType_c==pybullet.JOINT_PRISMATIC or jointType_c== pybullet.JOINT_REVOLUTE):
          jointIds_c.append(j)
print('转动or移动关节: ', jointIds_c)#这里输出的是转动关节数量，假设这里输出的关节索引是：[2, 3, 4, 5, 6, 7, 8,]，对你的正向运动学很关键
#下面的循环主要用来获取机械臂，这一步的作用是用来获取你的机械臂的末端索引号 
for joint_index in range(num_joints):
       joint_info = pybullet.getJointInfo(pybullet_robot, joint_index)
       print(f"\
               [0]关节索引:{joint_info[0]}\n\
               [1]关节名称:{joint_info[1]}\n\
               [2]关节类型:{joint_info[2]}\n\
               [3]此主体的位置状态变量中的第一个位置索引:{joint_info[3]}\n\
               [4]在这个物体的速度状态变量中的第一个速度索引:{joint_info[4]}\n\
               [5]保留参数:{joint_info[5]}\n\
               [6]关节阻尼大小:{joint_info[6]}\n\
               [7]关节摩擦系数:{joint_info[7]}\n\
               [8]滑动或旋转关节的位置下限:{joint_info[8]}\n\
               [9]滑动或旋转关节的位置上限:{joint_info[9]}\n\
               [10]关节最大力矩:{joint_info[10]}\n\
               [11]关节最大速度:{joint_info[11]}\n\
               [12]连杆名称:{joint_info[12]}\n\
               [13]在当前连杆坐标系中表示的移动或转动的关节轴(忽略JOINT_FIXED固定关节):{joint_info[13]}\n\
               [14]在父连杆坐标系中表示的关节位置:{joint_info[14]}\n\
               [15]在父连杆坐标系中表示的关节姿态(四元数x、y、z、w):{joint_info[15]}\n\
               [16]父连杆的索引,若是base连杆则返回-1:{joint_info[16]}\n\n")
# 利用上面循环假设获取关节索引号9是你的末端索引，假设关节名称end_eff_index，这个关节索引就是用来完成机械臂的逆向运动学重要条件，即：
end_eff_index = 6
#初始化关节角度

Position = [-0.0901079997420311, 0.7426530122756958, 0.19129900634288788]#目标位置，x，y，z
Orientation = [-1.5369999408721924, -0.004000000189989805, -1.5920000076293945, 0.020301342010498047]#目标姿态，这里是四元数形式作为输入
Joint_pos = pybullet.calculateInverseKinematics(pybullet_robot, end_eff_index, Position, Orientation, maxNumIterations=200) # 输入位置和姿态利用逆向运动学求解关节角度
    

for i in range(num_joints-1):#初始化关节位置为0
    pybullet.resetJointState(pybullet_robot, i, Joint_pos[i])
    pybullet.setJointMotorControl2(pybullet_robot, i, pybullet.POSITION_CONTROL, targetPosition=0.0)
    #下面给一个目标位姿来展示正逆向运动学的使用
    
#下面实现正向运动学来完成对机械臂的控制
pos = np.load('realdmp3d.npy')
jjj = np.load('realdmp3d_joint.npy')
joint_list = []
for i in range(pos.shape[0]):
    print(i)
    Joint_poss = pybullet.calculateInverseKinematics(pybullet_robot, end_eff_index, Position, Orientation, maxNumIterations=200) 
    for j in range(6):#初始化关节位置
        pybullet.resetJointState(pybullet_robot, j, Joint_poss[j])
        pybullet.setJointMotorControl2(pybullet_robot, j,  pybullet.POSITION_CONTROL, Joint_poss[j])
    pybullet.stepSimulation()#更新状态 每次控制关节角后都需要状态
    time.sleep(0.5)
    joint_list.append(Joint_pos)
joint_list = np.array(joint_list)
np.save('realdmp3d_joint.npy', joint_list)

    
